from enum import Enum
from typing import Dict, Optional

from core.agents.base import BaseAgent
from core.llm.factory import ModelType
from core.utils.log import mylogger


class AgentType(Enum):
    """Agent type enumeration."""

    DIALOGUE = "dialogue"
    MEMORY = "memory"
    EVAL = "eval"


def create_agent(
    agent_type: AgentType,
    name: str,
    model_type: ModelType = ModelType.DOUBAO_FLASH,
    **kwargs,
) -> BaseAgent:
    """Create agent instance following the factory pattern.

    Args:
        agent_type: Type of agent to create
        name: Agent name
        model_type: LLM model type
        **kwargs: Additional agent-specific parameters

    Returns:
        Agent instance
    """
    AGENT_MAPPINGS = {
        AgentType.DIALOGUE: lambda: __import_and_create("core.agents.dialogue_agent", "DialogueAgent", name, model_type, **kwargs),
        AgentType.MEMORY: lambda: __import_and_create("core.agents.memory_agent", "MemoryAgent", name, model_type, **kwargs),
        AgentType.EVAL: lambda: __import_and_create("core.agents.eval_agent", "EvalAgent", name, model_type, **kwargs),
    }
    
    if agent_type in AGENT_MAPPINGS:
        return AGENT_MAPPINGS[agent_type]()
    else:
        raise ValueError(f"Unsupported agent type: {agent_type}")


def __import_and_create(module_path: str, class_name: str, name: str, model_type: ModelType, **kwargs) -> BaseAgent:
    """Helper function to dynamically import and create agent instances."""
    module = __import__(module_path, fromlist=[class_name])
    agent_class = getattr(module, class_name)
    return agent_class(name=name, model_type=model_type, **kwargs)


class AgentManager:
    """Agent manager for creating and caching agent instances."""

    def __init__(self):
        self._agents: Dict[str, BaseAgent] = {}

    def get_agent(
        self,
        agent_type: AgentType,
        name: str,
        model_type: ModelType = ModelType.DOUBAO_FLASH,
        **kwargs,
    ) -> BaseAgent:
        """Get or create agent instance.

        Args:
            agent_type: Type of agent
            name: Agent name
            model_type: LLM model type
            **kwargs: Additional parameters

        Returns:
            Agent instance
        """
        agent_key = f"{agent_type.value}_{name}_{model_type.value}"

        if agent_key not in self._agents:
            self._agents[agent_key] = create_agent(
                agent_type=agent_type, name=name, model_type=model_type, **kwargs
            )
            mylogger.info(
                f"Created agent: {name} (type: {agent_type.value}, model: {model_type.value})"
            )

        return self._agents[agent_key]

    def get_dialogue_agent(
        self, name: str, model_type: ModelType = ModelType.DOUBAO_FLASH, **kwargs
    ) -> BaseAgent:
        """Convenience method for dialogue agents."""
        return self.get_agent(AgentType.DIALOGUE, name, model_type, **kwargs)

    def get_memory_agent(
        self, name: str, model_type: ModelType = ModelType.DOUBAO_FLASH, **kwargs
    ) -> BaseAgent:
        """Convenience method for memory agents."""
        return self.get_agent(AgentType.MEMORY, name, model_type, **kwargs)

    def get_eval_agent(
        self, name: str, model_type: ModelType = ModelType.DOUBAO_SEED, **kwargs
    ) -> BaseAgent:
        """Convenience method for evaluation agents."""
        return self.get_agent(AgentType.EVAL, name, model_type, **kwargs)


    def clear_agents(self, agent_type: Optional[AgentType] = None):
        """Clear cached agents.

        Args:
            agent_type: Specific agent type to clear, or None for all
        """
        if agent_type is None:
            self._agents.clear()
            mylogger.info("Cleared all cached agents")
        else:
            keys_to_remove = [
                key
                for key in self._agents.keys()
                if key.startswith(f"{agent_type.value}_")
            ]
            for key in keys_to_remove:
                del self._agents[key]
            mylogger.info(f"Cleared {agent_type.value} agents")


# Global agent manager instance
agent_manager = AgentManager()
